import com.sun.jna.Pointer;
import library.WintunLibrary;
import org.pcap4j.packet.*;
import org.pcap4j.packet.namednumber.DnsResourceRecordType;
import org.pcap4j.packet.namednumber.IpNumber;
import org.pcap4j.packet.namednumber.IpVersion;
import org.pcap4j.packet.namednumber.UdpPort;

import java.io.IOException;
import java.net.*;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

public class UdpData {

    static ConcurrentHashMap<String, UdpPipe> pipeMap = new ConcurrentHashMap<>();

    public static void parse(IpPacket packet, Pointer pointer) {

        //通过源ip和源端口来判断同一个udp流
        String key = getKey(packet);

        UdpPipe udpPipe = pipeMap.get(key);
        if (udpPipe == null) {
            //创建一个新的对象
            udpPipe = new UdpPipe(pointer);
            //防止多线程覆盖对象
            pipeMap.putIfAbsent(key, udpPipe);

        }
        udpPipe.receive(packet);

    }

    private static String getKey(IpPacket packet) {
        InetAddress srcAddr = packet.getHeader().getSrcAddr();
        UdpPacket udpPacket = (UdpPacket) packet.getPayload();
        UdpPort srcPort = udpPacket.getHeader().getSrcPort();
        return srcAddr + ":" + srcPort.valueAsInt();
    }

    private static void remove(IpPacket packet) {
        pipeMap.remove(getKey(packet));
    }

    static class UdpPipe {
        private Pointer pointer;

        private IpPacket srcPacket;

        private volatile DatagramChannel channel = null;

        private Thread readUdpThread;

        public UdpPipe(Pointer pointer) {
            this.pointer = pointer;
        }

        private void writeData(byte[] data) {

            Pointer packetPointer = WintunLibrary.INSTANCE.WintunAllocateSendPacket(this.pointer, data.length);

            if (packetPointer != null) {

                packetPointer.write(0, data, 0, data.length);

                // 发送数据包
                WintunLibrary.INSTANCE.WintunSendPacket(this.pointer, packetPointer);

            } else {
                System.err.println("Failed to allocate send packet.");
            }
        }

        public void receive(IpPacket ipPacket) {

            InetAddress dstAddr = ipPacket.getHeader().getDstAddr();
            UdpPacket udpPacket = (UdpPacket) ipPacket.getPayload();
            int dstPort = udpPacket.getHeader().getDstPort().valueAsInt();

            byte[] rawData = udpPacket.getPayload().getRawData();

            //过滤一些协议
            if (dstPort == 1900) {
                //ssdp协议
                remove(ipPacket);
                return;
            } else if (dstPort == 5355) {
                //LLMNR
                remove(ipPacket);
                return;
            } else if (dstPort == 5353) {
                //MDNS
                remove(ipPacket);
                return;
            } else if (dstAddr.getAddress()[3] == (byte) 255) {
                //广播
                remove(ipPacket);
                return;
            }else if(dstAddr.getAddress()[0]==(byte) 172 && dstAddr.getAddress()[1]==(byte) 29 ){
               if(dstAddr.getAddress()[2]!=(byte) 1 && dstAddr.getAddress()[2]!=(byte) 2){
                   remove(ipPacket);
                   return;
               }
            }

            //双重校验 volatile添加线程可见性
            if (channel == null) {
                synchronized (this) {
                    if (channel == null) {
                        try {
                            //判断是ip4还ip6
                            ProtocolFamily protocolFamily = StandardProtocolFamily.INET;
                            if (ipPacket instanceof IpV6Packet) {
                                protocolFamily = StandardProtocolFamily.INET6;
                            }

                            channel = DatagramChannel.open(protocolFamily);

                            if (dstAddr.getHostAddress().equals("172.29.1.2")) {
                                dstAddr = InetAddress.getByName("192.168.1.1");
                            }


                            channel.connect(new InetSocketAddress(dstAddr, dstPort));

                            srcPacket = ipPacket;

                            readUdpThread = new Thread(() -> {
                                ByteBuffer buffer = ByteBuffer.allocate(1024);

                                while (!Thread.interrupted()) {
                                    try {
                                        int len = channel.read(buffer);
                                        if (len > 0) {
                                            buffer.flip();
                                            byte[] array = new byte[buffer.limit()];
                                            buffer.get(array);

                                            //判断是dns 服务
                                            if (srcPacket.getHeader().getDstAddr().getHostAddress().equals("172.29.1.2")) {
                                                array = updateDnsIp(array);
                                            }

                                            //构建upd包
                                            byte[] bytes = buildUdpPacket(ipPacket, array);

                                            //写入网络适配器
                                            writeData(bytes);

                                            //释放连接
                                            close();

                                        } else {
                                            //认为连接关闭
                                            close();
                                        }

                                    } catch (Exception e) {
                                        e.printStackTrace();
                                        close();
                                    }
                                }

                            });

                            readUdpThread.setName(dstAddr + ":" + dstPort + " readUdpThread");
                            readUdpThread.start();

                        } catch (Exception e) {
                            e.printStackTrace();
                        }

                    }
                }
            }
            try {
                if (channel.isConnected()) {
                    //写入数据
                    channel.write(ByteBuffer.wrap(rawData));
                }

            } catch (IOException e) {
                e.printStackTrace();
                close();
            }

        }

        private void close() {
            if (readUdpThread != null) {
                readUdpThread.interrupt();
            }
            if (channel != null) {
                try {
                    channel.close();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            remove(srcPacket);
        }

        private byte[] updateDnsIp(byte[] array) throws IllegalRawDataException {

            DnsPacket dnsPacket = DnsPacket.newPacket(array, 0, array.length);
            DnsPacket.DnsHeader header = dnsPacket.getHeader();
            String domain = null;
            if (header.getQdCount() > 0) {
                DnsQuestion dnsQuestion = header.getQuestions().get(0);
                domain = dnsQuestion.getQName().getName();
                System.out.println(domain);
            }


            //（附加信息）部分
            List<DnsResourceRecord> answersList = new ArrayList<>();
            for (DnsResourceRecord answer : header.getAnswers()) {
                if (answer.getDataType() == DnsResourceRecordType.A) {

                    DnsRDataA dnsRDataA = (DnsRDataA) answer.getRData();

                    byte[] bytes = IpCache.putCache(dnsRDataA.getAddress());
                    DnsRDataA dnsRDataA1 = DnsRDataA.newInstance(bytes, 0, bytes.length);


                    DnsResourceRecord build = new DnsResourceRecord.Builder()
                            .name(answer.getName())
                            .dataType(answer.getDataType())
                            .ttl(answer.getTtl())
                            .rData(dnsRDataA1)
                            .dataClass(answer.getDataClass())
                            .correctLengthAtBuild(true).build();
                    answersList.add(build);

                } else {
                    answersList.add(answer);
                }

            }
            //（附加信息）部分
            List<DnsResourceRecord> additionalList = new ArrayList<>();
            for (DnsResourceRecord info : header.getAdditionalInfo()) {
                if (info.getDataType() == DnsResourceRecordType.A) {

                    DnsRDataA dnsRDataA = (DnsRDataA) info.getRData();

                    byte[] bytes = IpCache.putCache(dnsRDataA.getAddress());
                    DnsRDataA dnsRDataA1 = DnsRDataA.newInstance(bytes, 0, bytes.length);


                    DnsResourceRecord build = new DnsResourceRecord.Builder()
                            .name(info.getName())
                            .dataType(info.getDataType())
                            .ttl(info.getTtl())
                            .rData(dnsRDataA1)
                            .dataClass(info.getDataClass())
                            .correctLengthAtBuild(true).build();
                    additionalList.add(build);

                } else {
                    additionalList.add(info);
                }
            }


            DnsPacket build = new DnsPacket.Builder()
                    .id(dnsPacket.getHeader().getId())
                    .response(dnsPacket.getHeader().isResponse())
                    .opCode(dnsPacket.getHeader().getOpCode())
                    .authoritativeAnswer(dnsPacket.getHeader().isAuthoritativeAnswer())
                    .truncated(dnsPacket.getHeader().isTruncated())
                    .recursionDesired(dnsPacket.getHeader().isRecursionDesired())
                    .recursionAvailable(dnsPacket.getHeader().isRecursionAvailable())
                    .reserved(dnsPacket.getHeader().getReservedBit())
                    .authenticData(dnsPacket.getHeader().isAuthenticData())
                    .checkingDisabled(dnsPacket.getHeader().isCheckingDisabled())
                    .rCode(dnsPacket.getHeader().getrCode())
                    .questions(dnsPacket.getHeader().getQuestions())
                    .answers(answersList)
                    .authorities(dnsPacket.getHeader().getAuthorities())
                    .additionalInfo(additionalList)
                    .qdCount(dnsPacket.getHeader().getQdCount())
                    .anCount(dnsPacket.getHeader().getAnCount())
                    .nsCount(dnsPacket.getHeader().getNsCount())
                    .arCount(dnsPacket.getHeader().getArCount()).build();


            return build.getRawData();
        }

        private byte[] buildUdpPacket(IpPacket ipPacket, byte[] data) {

            UdpPacket udpPacket = (UdpPacket) ipPacket.getPayload();

            UdpPort srcPort = udpPacket.getHeader().getSrcPort();
            UdpPort dstPort = udpPacket.getHeader().getDstPort();


            UnknownPacket.Builder unknownBuilder = new UnknownPacket.Builder();
            unknownBuilder.rawData(data);

            //创建一个udp包
            UdpPacket.Builder udpBuilder = new UdpPacket.Builder()
                    .srcPort(dstPort)
                    .dstPort(srcPort)
                    .payloadBuilder(unknownBuilder)
                    .correctLengthAtBuild(true);


            InetAddress srcAddr = ipPacket.getHeader().getSrcAddr();
            InetAddress dstAddr = ipPacket.getHeader().getDstAddr();

            //创建一个ip包
            IpPacket newIpPacket;
            if (ipPacket instanceof IpV6Packet) {
                newIpPacket = new IpV6Packet.Builder()
                        .version(IpVersion.IPV6)
                        .srcAddr((Inet6Address) dstAddr)
                        .dstAddr((Inet6Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .payloadBuilder(udpBuilder)
                        .build();
            } else {

                newIpPacket = new IpV4Packet.Builder()
                        .version(IpVersion.IPV4)
                        .ttl((byte) 45)
                        .tos(IpV4Rfc1349Tos.newInstance((byte) 0))
                        .protocol(IpNumber.UDP)
                        .srcAddr((Inet4Address) dstAddr)
                        .dstAddr((Inet4Address) srcAddr)
                        .correctLengthAtBuild(true)
                        .correctChecksumAtBuild(true)
                        .payloadBuilder(udpBuilder)
                        .identification((byte) TcpData.getIpId())
                        .build();
            }
            return newIpPacket.getRawData();
        }
    }


}
